# -*- coding: utf-8 -*-
'''
Created on 06.07.2020

@author: yu03
'''
import imageio
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
import glob
import os
import shutil
from matplotlib.gridspec import GridSpec
import matplotlib.patches as patches
from plt_style import *
colors = color_set_medium

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]

pix_size = 5.3e-6
pix_num = 200
screen_diameter = pix_num * pix_size


dx = np.linspace(0, (pix_num-1)*pix_size, num=pix_num)
dy = dx
V_r_x, V_r_y, V_r_z = -0.002, -0.002, 1 # Reference vector
L, D = 0.1, 0.1  # D = displacement, L is static
V_m_x, V_m_y, V_z_m = 0, 0.000, 1 # Measurement vector

X, Y = np.meshgrid(dx, dy)

def f(X, Y):
    '''
        Window 2d
    '''
    window = signal.gaussian(pix_num, std=pix_num/8)
    window_x, window_y = np.meshgrid(window, window)
    window_2d = window_x*window_y
    
    k = 2 * np.pi / Lamda
    I_0 = 512
    
    '''
        Wave Generating
    '''
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
#     d_ref = X*2*V_r_x*(1-V_r_x**2-V_r_y**2)/(1-V_r_x**2)/(1+V_r_x**2+V_r_y**2) + Y*2*V_r_y*(1-V_r_x**2-V_r_y**2)/(1-V_r_y**2)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
#     d_mea = X*2*V_m_x*(1-V_m_x**2-V_m_y**2)/(1-V_m_x**2)/(1+V_m_x**2+V_m_y**2) + Y*2*V_m_y*(1-V_m_x**2-V_m_y**2)/(1-V_m_y**2)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    diff_phi = k * (diff_Z_p + (d_mea-d_ref))
    I_beat = I_0*(1+np.cos(diff_phi))
#     I_beat = I_beat.astype(np.int)
    
    '''
        Wave Generating (2nd way)
    '''
    D_d = D + X * (V_m_x - V_r_x) + Y * (V_m_y - V_r_y)
    L_d = L + X * V_r_x + Y * V_r_y
    diff_L_d_2 = D_d + (L_d + D_d) * (1 - (V_m_x**2 + V_m_y**2)) / (1 + (V_m_x**2 + V_m_y**2)) - L_d * (1 - (V_r_x**2 + V_r_y**2)) / (1 + (V_r_x**2 + V_r_y**2))
    Z = I_0 * (1 + np.cos(2 * np.pi * diff_L_d_2 / Lamda))
    
    return (I_beat) * window_2d

path = r'C:\Users\yu03\Desktop\tilted'
images = []
'''
    Generating image set
'''
# if 1:
#     fig = plt.figure()
#     fig.set_size_inches(10, 3, True)
#     gs = GridSpec(1, 2, width_ratios=[3, 1])
#     ax1 = fig.add_subplot(gs[0])
#     ax2 = fig.add_subplot(gs[1])
    
#     laser = ax1.add_patch(patches.Rectangle((5, 50),15,10,facecolor=colors['blue'],fill=True))
#     BS = ax1.plot([30,45],[41,66],linewidth=1,color=colors['blue'])
#     Ref_M = ax1.plot([30,45],[90,85],linewidth=3,color=colors['blue'],zorder=100)
#     Mea_M, = ax1.plot([60,60],[40,65],linewidth=3,color=colors['blue'])
#     Camera = ax1.add_patch(patches.Rectangle((30, 10),15,10,facecolor=colors['blue'],fill=True))
#     ax1.plot([38,38],[20,87],color=colors['red'])
#     light, = ax1.plot([20,60],[55,55],color=colors['red'])
#     
#     ax1.text(0.08, 0.65, '激光器', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
#     ax1.text(0.33, 0.95, '参考镜', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
#     ax1.text(0.33, 0.05, '探测器', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
#     ax1.text(0.7, 0.3, '被测目标', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
#     
#     ax1.set_xlim(0,100)
#     ax1.set_ylim(0,100)
#     ax1.spines['bottom'].set_visible(False)
#     ax1.spines['left'].set_visible(False)
#     ax1.tick_params(axis='both',which='both',labelbottom=[],labelleft=[],length=0)

def layout_plot(M_x1,M_x2,light_x2,text,num,unit):
    fig = plt.figure()
    fig.set_size_inches(10, 3, True)
    gs = GridSpec(1, 2, width_ratios=[2, 1])
    ax1 = fig.add_subplot(gs[0])
    ax2 = fig.add_subplot(gs[1])
    img = ax2.imshow(f(X, Y),cmap='gray')
    img.set_clim([0,1000])
    line, = ax2.plot((-f(X, Y)[pix_num//2]/(5000//pix_num)+pix_num-10), color='white', linewidth=0.5)
    line_ver, = ax2.plot(f(X, Y)[:,pix_num//2]/(5000//pix_num)+10, np.arange(pix_num), color='white', linewidth=0.5)
    ax2.axis('off')
    
    laser = ax1.add_patch(patches.Rectangle((5, 50),15,10,facecolor=colors['blue'],fill=True))
    BS = ax1.plot([30,45],[41,66],linewidth=1,color=colors['blue'])
    Ref_M = ax1.plot([30,45],[90,85],linewidth=3,color=colors['blue'],zorder=100)
    Mea_M, = ax1.plot([M_x1,M_x2],[40,65],linewidth=3,color=colors['blue'])
    Camera = ax1.add_patch(patches.Rectangle((30, 10),15,10,facecolor=colors['blue'],fill=True))
    ax1.plot([38,38],[20,87],color=colors['red'])
    light, = ax1.plot([20,light_x2-0.5],[55,55],color=colors['red'])
    
    ax1.text(0.08, 0.65, '激光器', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
    ax1.text(0.33, 0.95, '参考镜', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
    ax1.text(0.33, 0.05, '探测器', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
    ax1.text((M_x1-10)/100, 0.3, '被测目标', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)

    ax1.set_xlim(0,100)
    ax1.set_ylim(0,100)
    ax1.spines['bottom'].set_visible(False)
    ax1.spines['left'].set_visible(False)
    ax1.tick_params(axis='both',which='both',labelbottom=[],labelleft=[],length=0)
    ax1.text(0.65, 0.75, text, fontsize=15, horizontalalignment='left', verticalalignment='center', color=color_set_normal['red'], transform=ax1.transAxes)
    
    if num>unit*4 and num < unit*5:
#         style = "Simple, tail_width=0.5, head_width=4, head_length=8"
#         kw = dict(arrowstyle=style, color="k")
#         arrow = ax1.add_patch(patches.FancyArrowPatch((85, 50), (95, 50), connectionstyle="arc3,rad=.5", **kw))
        ax1.plot([60],[50],color=colors['black'],marker=r'$\circlearrowleft$',ms=30)
#         plt.show()
    if num>unit*5 and num < unit*6:
        ax1.plot([60],[50],color=colors['black'],marker=r'$\circlearrowright$',ms=30)
    return
    
    
if 1:
    unit=50
    for num in range(int(unit*6)):
        print(num)
        if num<unit:
            D += -Lamda / 16
            D_layout = (0.1-D)/Lamda*10+60
            M_x1, M_x2 = D_layout, D_layout
            light_x2 = D_layout
            text = '位移'
            
        if num >= unit*1 and num < unit*2:
            V_m_y += -.8e-4
            D_layout = (0.1-D)/Lamda*10+60
            D_4_V = V_m_y*1e3    
            M_x1, M_x2 = D_layout-D_4_V, D_layout+D_4_V
            text = '旋转（俯仰）'
        
        if num >= unit*2 and num < unit*3:
            V_m_y += .8e-4
            D_layout = (0.1-D)/Lamda*10+60
            D_4_V = V_m_y*1e3    
            M_x1, M_x2 = D_layout-D_4_V, D_layout+D_4_V
            text = '旋转（俯仰）'

        if num >=unit*3 and num <unit*4:
            D += Lamda / 16
            D_layout = (0.1-D)/Lamda*10+60
            D_4_V = V_m_y*1e3
            M_x1, M_x2 = D_layout-D_4_V,D_layout+D_4_V
            light_x2 = D_layout
            text = '位移'
            
        if num >= unit*4 and num < unit*5:
            V_m_x += -.8e-4
            text = '旋转（偏摆）'
         
        if num >= unit*5 and num < unit*6:
            V_m_x += .8e-4
            text = '旋转（偏摆）'
         
        layout_plot(M_x1=M_x1,M_x2=M_x2,light_x2=light_x2,text=text, num=num,unit=unit)
        plt.savefig(path+'\\'+'%s.jpg'%num)
        plt.clf()
        plt.close()

'''
    Reading images
'''   
if 1:
    files = sorted(glob.glob(path+r'\*.jpg'), key=os.path.getmtime)
    for file in files: #对路径下的每个.pdf文件进行操作
#         print(file)
        img = imageio.imread(file)
        images.append(img)
    print(len(images))
    imageio.mimsave(r'C:\Users\yu03\Desktop\tilted.gif', images, duration=0.1)